"""
MCP 安全管理
实现 MCP 协议的安全机制,包括认证、授权和加密
"""

import base64
import hashlib
import hmac
import json
import logging
import secrets
from datetime import datetime, timedelta
from typing import Dict, Optional, Tuple

import jwt
from cryptography.fernet import Fernet

class MCPSecurity:
    """MCP 安全管理器
    
    提供 MCP 协议的安全功能:
    - 令牌生成和验证
    - 请求签名和验证
    - 数据加密和解密
    - 访问控制
    """
    
    def __init__(self, config: Dict):
        """初始化安全管理器
        
        Args:
            config: 安全配置
                secret_key: 密钥
                token_expiry: 令牌过期时间(秒)
                encryption_key: 加密密钥
                auth_method: 认证方式 (token/signature)
        """
        self.config = config
        self.logger = logging.getLogger(__name__)
        
        # 初始化加密器
        if "encryption_key" in config:
            self.fernet = Fernet(config["encryption_key"].encode())
        else:
            key = Fernet.generate_key()
            self.fernet = Fernet(key)
            config["encryption_key"] = key.decode()
            
    def generate_token(self, client_id: str, scopes: list) -> str:
        """生成访问令牌
        
        Args:
            client_id: 客户端标识
            scopes: 权限范围
            
        Returns:
            str: JWT 令牌
        """
        now = datetime.utcnow()
        expiry = now + timedelta(seconds=self.config.get("token_expiry", 3600))
        
        payload = {
            "client_id": client_id,
            "scopes": scopes,
            "iat": now,
            "exp": expiry
        }
        
        return jwt.encode(
            payload,
            self.config["secret_key"],
            algorithm="HS256"
        )
        
    def validate_token(self, token: str) -> Tuple[bool, Optional[Dict]]:
        """验证访问令牌
        
        Args:
            token: JWT 令牌
            
        Returns:
            Tuple[bool, Optional[Dict]]: (是否有效, 令牌数据)
        """
        try:
            payload = jwt.decode(
                token,
                self.config["secret_key"],
                algorithms=["HS256"]
            )
            return True, payload
        except jwt.ExpiredSignatureError:
            self.logger.warning("Token expired")
            return False, None
        except jwt.InvalidTokenError as e:
            self.logger.error(f"Invalid token: {e}")
            return False, None
            
    def sign_request(self, request: Dict) -> Dict:
        """签名请求
        
        Args:
            request: 请求数据
            
        Returns:
            Dict: 带签名的请求
        """
        # 生成随机数和时间戳
        nonce = secrets.token_hex(16)
        timestamp = str(int(datetime.utcnow().timestamp()))
        
        # 准备签名数据
        data = {
            "request": request,
            "nonce": nonce,
            "timestamp": timestamp
        }
        
        # 计算签名
        message = json.dumps(data, sort_keys=True).encode()
        signature = hmac.new(
            self.config["secret_key"].encode(),
            message,
            hashlib.sha256
        ).hexdigest()
        
        # 添加签名信息
        data["signature"] = signature
        return data
        
    def verify_signature(self, signed_request: Dict) -> Tuple[bool, Optional[Dict]]:
        """验证请求签名
        
        Args:
            signed_request: 带签名的请求
            
        Returns:
            Tuple[bool, Optional[Dict]]: (签名是否有效, 原始请求)
        """
        try:
            # 提取签名和数据
            signature = signed_request.pop("signature")
            
            # 验证时间戳
            timestamp = int(signed_request["timestamp"])
            now = int(datetime.utcnow().timestamp())
            if abs(now - timestamp) > 300:  # 5分钟有效期
                return False, None
                
            # 计算签名
            message = json.dumps(signed_request, sort_keys=True).encode()
            expected_signature = hmac.new(
                self.config["secret_key"].encode(),
                message,
                hashlib.sha256
            ).hexdigest()
            
            # 验证签名
            if not hmac.compare_digest(signature, expected_signature):
                return False, None
                
            return True, signed_request["request"]
            
        except (KeyError, ValueError) as e:
            self.logger.error(f"Invalid signed request: {e}")
            return False, None
            
    def encrypt_data(self, data: str) -> str:
        """加密数据
        
        Args:
            data: 原始数据
            
        Returns:
            str: 加密后的数据
        """
        return self.fernet.encrypt(data.encode()).decode()
        
    def decrypt_data(self, encrypted_data: str) -> str:
        """解密数据
        
        Args:
            encrypted_data: 加密数据
            
        Returns:
            str: 解密后的数据
        """
        return self.fernet.decrypt(encrypted_data.encode()).decode()
        
    def check_permission(self, token_payload: Dict, required_scopes: list) -> bool:
        """检查权限
        
        Args:
            token_payload: 令牌数据
            required_scopes: 所需权限
            
        Returns:
            bool: 是否有权限
        """
        if not token_payload or "scopes" not in token_payload:
            return False
            
        token_scopes = set(token_payload["scopes"])
        required_scopes = set(required_scopes)
        
        return required_scopes.issubset(token_scopes)
        
    def audit_log(self, event: str, details: Dict) -> None:
        """记录审计日志
        
        Args:
            event: 事件类型
            details: 事件详情
        """
        log_entry = {
            "timestamp": datetime.utcnow().isoformat(),
            "event": event,
            "details": details
        }
        
        self.logger.info(f"Audit: {json.dumps(log_entry)}")
        
    def rotate_keys(self) -> None:
        """轮换密钥"""
        # 生成新的加密密钥
        new_key = Fernet.generate_key()
        old_fernet = self.fernet
        self.fernet = Fernet(new_key)
        self.config["encryption_key"] = new_key.decode()
        
        # 生成新的签名密钥
        new_secret = secrets.token_hex(32)
        old_secret = self.config["secret_key"]
        self.config["secret_key"] = new_secret
        
        self.logger.info("Security keys rotated successfully")
        
        # TODO: 实现密钥过渡期处理 